package handler;

import java.io.BufferedReader;
import java.io.DataOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.ArrayList;
import java.util.List;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;
import javax.crypto.SecretKeyFactory;
import javax.crypto.spec.DESKeySpec;

/**
 * Handles the encryption and decryption of streams. This class is
 * essentially a utility class that is used elsewhere in the program to support
 * encrypted Notebook files. Notebook entries are encrypted with the DES
 * algorithm. Helper methods are also provided for performing RSA encryption.
 * RSA is used in the cloud feature of the Math Graph program. Each user signs
 * the primary key of the files they have access to. None of the methods in this
 * class will close the InputStream/OutputStream they are provided with. It is
 * the responsibility of the caller to close these provided streams.
 * 
 * @author Jared Schlicht
 */
public class EncryptionHandler {
    /**
     * Specifies the number of bits to use when creating RSA keys.
     */
    private static final int KEY_SIZE = 512;
    
    /**
     * Uses the specified password to encrypt data from the specified input
     * stream. The encrypted results are sent to the specified output stream.
     * The caller is responsible for closing both the input and output streams.
     * 
     * @author Jared Schlicht
     * @param password the password to use to encrypt the data
     * @param inputStream the source of data that should be encrypted
     * @param outputStream the result of encrypting the data from the input stream
     * @return true if the data was encrypted without errors, false otherwise
     */
    public boolean encryptDES(String password, InputStream inputStream, OutputStream outputStream) {
        return this.performDESOperation(password, inputStream, outputStream, Cipher.ENCRYPT_MODE);
    }
    
    /**
     * Encrypts the specified input stream using the provided public key.
     * The encrypted data is sent to the provided output stream. The caller
     * is responsible for closing both the input and output streams.
     * 
     * @author Jared Schlicht
     * @param publicKey the public key to use to encrypt the specified input stream
     * @param inputStream the source of data that should be RSA encrypted
     * @param outputStream the result of encrypting the data from the input stream
     * @return true if the data was encrypted without errors, false otherwise
     */
    public boolean encryptRSA(byte[] publicKey, InputStream inputStream, OutputStream outputStream) {
        try {
            KeyFactory keyFactory = KeyFactory.getInstance("rsa");
            X509EncodedKeySpec publicKeySpec = new X509EncodedKeySpec(publicKey);
            PublicKey rsaPublicKey = keyFactory.generatePublic(publicKeySpec);
            
            return this.performRSAOperation(rsaPublicKey, inputStream, outputStream, Cipher.ENCRYPT_MODE);
            
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
        } catch (InvalidKeySpecException e) {
            e.printStackTrace();
        }
        
        return false;
    }
    
    /**
     * Decrypts the specified input stream using the provided private key.
     * The decrypted data is sent to the provided output stream. The caller
     * is responsible for closing the provided input and output streams.
     * 
     * @author Jared Schlicht
     * @param privateKey the private key to use to decrypt the specified input stream
     * @param inputStream the source of RSA-encrypted data that should be decrypted
     * @param outputStream the result of decrypting the data from the input stream
     * @return true if the data was decrypted without errors, false otherwise
     */
    public boolean decryptRSA(byte[] privateKey, InputStream inputStream, OutputStream outputStream) {
        try {
            KeyFactory keyFactory = KeyFactory.getInstance("rsa");
            
            PKCS8EncodedKeySpec privateKeySpec = new PKCS8EncodedKeySpec(privateKey);
            PrivateKey rsaPrivateKey = keyFactory.generatePrivate(privateKeySpec);
            
            return this.performRSAOperation(rsaPrivateKey, inputStream, outputStream, Cipher.DECRYPT_MODE);
            
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
        } catch (InvalidKeySpecException e) {
            e.printStackTrace();
        }
        
        return false;
    }
    
    /**
     * Helper method for performing RSA encryption or decryption. The caller
     * is responsible for closing the provided streams.
     * 
     * @author Jared Schlicht
     * @param key the key to use to encrypt/decrypt the data. This should be
     * a public key for encryption and a private key for decryption
     * @param inputStream data will be read from this source
     * @param outputStream the result of the rsa operation will be sent to this stream
     * @param opMode Specifies whether encryption or decryption should be
     * performe.d Cipher.DECRYPT_MODE will attempt to decrypt the input stream.
     * Cipher.ENCRYPT_MODE will attempt to encrypt the input stream.
     * @return true if the RSA operation was completed without errors, false otherwise
     */
    private boolean performRSAOperation(Key key, InputStream inputStream, OutputStream outputStream, int opMode) {
        try {
            Cipher rsa = Cipher.getInstance("rsa");
            rsa.init(opMode, key);
            
            byte[] inputBuffer = new byte[64];
            while (true) {
                int bytesRead = inputStream.read(inputBuffer);
                if(bytesRead == -1) {
                    break;
                }
                
                byte[] output = rsa.update(inputBuffer, 0, bytesRead);
                if (output != null) {
                    outputStream.write(output);
                }
            }
            
            byte[] output = rsa.doFinal();
            if (output != null) {
                outputStream.write(output);
            }
            outputStream.flush();
            return true;
        } catch (InvalidKeyException e) {
            e.printStackTrace();
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
        } catch (NoSuchPaddingException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        } catch (IllegalBlockSizeException e) {
            e.printStackTrace();
        } catch (BadPaddingException e) {
            e.printStackTrace();
        }
        return false;
    }

    /**
     * Generates a public/private RSA key pair using the specified password
     * 
     * @author Jared Schlicht
     * @param password the password to use to generate the public and private keys
     * @return A KeyPair object (containing both the public and private keys) or
     * null if the RSA algorithm could not be loaded.
     */
    public KeyPair createKeys(String password) {
        try {
            KeyPairGenerator kg = KeyPairGenerator.getInstance("rsa");
            SecureRandom secureSeed = SecureRandom.getInstance("SHA1PRNG");
            secureSeed.setSeed(password.getBytes());
            kg.initialize(KEY_SIZE, secureSeed);
            KeyPair pair = kg.generateKeyPair();
            return pair;
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
        }
        
        return null;
    }
    
    /**
     * A helper method to encrypt or decrypt DES data from the input stream.
     * The result is sent it to the specified output stream.
     * 
     * @author Jared Schlicht
     * @param password the password to use to encrypt/decrypt
     * @param inputStream stream to read data from
     * @param outputStream stream to output data to
     * @param opMode Specifies whether encryption or decryption should be
     * performed. Cipher.ENCRYPT_MODE to encrypt data, Cipher.DECRYPT_MODE to
     * decrypt data
     * @return true if the operation completed without errors, false otherwise
     */
    private boolean performDESOperation(String password, InputStream inputStream, OutputStream outputStream, int opMode) {
        /* DES passwords must be at least 8 characters. Append zeroes until the
         * password is at least 8 characters. Note that this is terribly
         * insecure and should never be performed in production code.
         */
        while (password.length() < 8)
        {
            password = password + "0";
        }
        
        byte[] desKeyData = password.getBytes();
        try {
            DESKeySpec desKeySpec = new DESKeySpec(desKeyData);
            SecretKeyFactory keyFactory = SecretKeyFactory.getInstance("DES");
            SecretKey desKey = keyFactory.generateSecret(desKeySpec);
            Cipher des = Cipher.getInstance("DES");
            des.init(opMode, desKey);
            
            byte[] inputBuffer = new byte[64];
            while (true) {
                int bytesRead = inputStream.read(inputBuffer);
                if(bytesRead == -1) {
                    break;
                }
                
                byte[] output = des.update(inputBuffer, 0, bytesRead);
                if (output != null) {
                    outputStream.write(output);
                }
            }
            
            byte[] output = des.doFinal();
            if (output != null) {
                outputStream.write(output);
            }
            outputStream.flush();
            return true;
        } catch (InvalidKeyException e) {
            e.printStackTrace();
        } catch (NoSuchAlgorithmException e) {
            e.printStackTrace();
        } catch (InvalidKeySpecException e) {
            e.printStackTrace();
        } catch (NoSuchPaddingException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        } catch (IllegalBlockSizeException e) {
            e.printStackTrace();
        } catch (BadPaddingException e) {
            e.printStackTrace();
        }
        return false;
    }
    
    /**
     * Uses the specified password to decrypt data from the specified input
     * stream. The decrypted results are sent to the specified output stream.
     * The caller is responsible for closing both the input and output streams.
     * 
     * @author Jared Schlicht
     * @param password the password to use to encrypt the data
     * @param inputStream the source of encrypted data that should be decrypted
     * @param outputStream the result of decrypting the data from the input stream
     * @return true if the data was decrypted without errors, false otherwise
     */
    public boolean decryptDES(String password, InputStream inputStream, OutputStream outputStream) {
        return this.performDESOperation(password, inputStream, outputStream, Cipher.DECRYPT_MODE);
    }
}
